"""
(Possibly marginally secure) RPC server

Packet is:

    octets    meaning
    4:          seclen = length of secret
    <seclen>:   secret
    4:          msglen = length of message
    <msglen>:   message (code)

integers are big-endian, 32 bits.


 A minimal windows service

 Based on _Python Programming on win32_, pp347+

 install by calling with argument "install"

"""

from __future__ import absolute_import
from __future__ import with_statement

__metaclass__ = type

import os
import sys
import pickle
import select
import logging
import threading

from SocketServer import (TCPServer,
                          StreamRequestHandler)
from cStringIO import StringIO

logging.basicConfig()
log = logging.getLogger(__name__)

try:
    import win32serviceutil
    import win32service
    import win32event
    HAVE_WIN32 = True
except ImportError:
    HAVE_WIN32 = False


class Packet:

    class Error(Exception):
        pass

    def __init__(self, secret, msg):
        self.secret = secret
        self.msg = msg
        self.seclen = len(str(secret))
        self.msglen = len(str(msg))
        self.bytes = None

    def format(self):
        import struct
        return "%s%s%s%s" % (struct.pack(">I", self.seclen),
                             self.secret,
                             struct.pack(">I", self.msglen),
                             self.msg)
    def __str__(self):
        if not self.bytes:
            self.bytes = self.format()
        return self.bytes

    def __len__(self):
        return len(str(self))

    @staticmethod
    def read(secret, stream):
        import struct
        data = stream.read(4)
        seclen = struct.unpack(">I", data)[0]
        if seclen != len(secret):
            raise Packet.Error("Incorrect secret size")

        data = stream.read(seclen)
        if secret != data:
            raise Packet.Error("Incorrect secret")

        data = stream.read(4)

        msglen = struct.unpack(">I", data)[0]
        msg = stream.read(msglen)
        return msg


class Client:
    def __init__(self, secret, server, port=1666):
        self.secret = secret
        self.seclen = len(str(secret))
        self.server = server
        self.port = port

    def connect(self):
        import socket
        s = socket.socket()
        s.connect((self.server, self.port))
        return s

    def send(self, msg):
        socket = self.connect()
        msglen = len(str(msg))
        packet = Packet(self.secret, msg)
        try:
            sent = send_packet(socket, packet)
            if sent != len(packet):
                raise Exception("Failed to send whole packet")

            out = self.recv(socket)
        finally:
            socket.close()
        return out

    def call(self, func, *args):
        import inspect
        f = func.__name__
        args = ','.join(map(repr, args))
        source = inspect.getsource(func)
        source += "\n\n__return__ = %s(%s)\n" % (f, args)
        result = self.send(source)
        if isinstance(result, error):
            raise result
        return result

    def recv(self, socket):
        buf = []
        bufsize = 1024
        while True:
            data = socket.recv(bufsize)
            if data:
                buf.append(data)
            else:
                break
        return pickle.loads(''.join(buf))


def send_packet(socket, data):
    total = 0
    data = str(data)
    while True:
        sent = socket.send(data)
        total += sent
        if sent < len(data):
            data = data[total:]
        else:
            break
    return total

real_stdout = sys.stdout
real_stdin  = sys.stdin
real_stderr = sys.stderr
real_exit = os._exit

def dummy_exit(status):
    log_error("Attempt to call os._exit()")

# error classes for pickling
class error(Exception): pass
class winerror(error): pass
class strerror(error): pass

def execute(msg, stdin):
    sys.stdin  = stdin
    sys.stdout = StringIO()
    sys.stderr = StringIO()
    os._exit = dummy_exit

    env = {'__return__': None}
    ret = None
    try:
        exec msg in env
        ret = env['__return__']
    except:
        # global function, set below
        ret = handle_error(*sys.exc_info())
    finally:
        sys.stdin = real_stdin
        sys.stdout = real_stdout
        sys.stderr = real_stderr
        os._exit = real_exit


    try:
        ret = pickle.dumps(ret)
    except:
        ret = strerror(map(str, sys.exc_info()[:2]))
        ret = pickle.dumps(ret)
    return ret


class Handler(StreamRequestHandler):
    def handle(self):
        try:
            msg = Packet.read(self.server.secret, self.rfile)
        except Packet.Error:
            log_error("Error reading packet")
            self.connection.close()
            return
        ret = execute(msg, self.rfile)
        self.wfile.write(ret)


class Server(TCPServer):
    allow_reuse_address = True
    def verify_request(self, request, client):
        return client[0] in self.allowed_clients


def select_server(event, secret, address, clients):
    def run(event=event):
        server = make_server(secret, address, clients)
        socket = server.fileno()

        while event.isSet():
            ready = select.select([socket], [], [], 1)
            if socket in ready[0]:
                server.handle_request()
    return threading.Thread(target=run)


def make_server(secret, addr, clients):
    addr, port = addr.split(':')
    port = int(port)
    s = Server((addr, port), Handler)
    s.secret = secret
    s.seclen = len(secret)
    s.allowed_clients = tuple(clients)
    return s


def gen_secret(length):
    import string, random
    chars = string.letters + string.digits
    return ''.join(random.sample(chars*100, length))


def test_server(secret=None):
    if secret is None:
        secret = gen_secret(20)
        print "secret:", secret
    allowed = ['127.0.0.1','192.168.0.1']
    serv = make_server(secret, '0.0.0.0:6666', allowed)
    serv.serve_forever()


def test_client(secret, server=None, port=1666):
    cli = Client(secret, server or '', port)
    msg = "import string; __return__ = string.printable"
    result = cli.send(msg)
    import string
    assert result == string.printable


def test_good(*args):
    return "GOOD! " + str(args)

def test_exc(*args):
    return None + 1 + args

def test_exit(*args):
    import os
    os._exit(1)

def _tb(args):
    import traceback
    tb = traceback.format_exc()
    return "\n\n".join(list(args) + [tb])

def handle_error_str(E, e, tb):
    log_error(str(e))
    return strerror((str(E), str(e)))
handle_error = handle_error_str


if HAVE_WIN32:
    def log_win32(*args):
        import servicemanager
        servicemanager.LogErrorMsg(_tb(args))
    log_error = log_win32

    def handle_error_win32(E, e, tb):
        import pythoncom
        if isinstance(e, pythoncom.com_error):
            return winerror(tuple(e))
        return handle_error_str(E, e, tb)
    handle_error = handle_error_win32

    class AppService(win32serviceutil.ServiceFramework):
        _svc_name_ = "MyAppAgent"
        _svc_display_name_ = "My App Windows agent"

        def __init__(self, args):
            win32serviceutil.ServiceFramework.__init__(self, args)
            self.hWaitStop = win32event.CreateEvent(None, 0, 0, None)

        def get_config(self):
            dir = os.path.dirname(os.path.abspath(__file__))
            import ConfigParser
            config = ConfigParser.ConfigParser()
            config.read(os.path.join(dir, 'winagent.ini'))
            secret = config.get('myapp', 'secret')
            listen = config.get('myapp', 'listen')
            clients = config.get('myapp', 'clients').split()
            return secret, listen, clients

        def SvcStop(self):
            self.ReportServiceStatus(win32service.SERVICE_STOP_PENDING)
            win32event.SetEvent(self.hWaitStop)

        def SvcDoRun(self):
            event = threading.Event()
            event.set()
            secret, listen, clients = self.get_config()
            thread = select_server(event, secret, listen, clients)
            thread.start()

            win32event.WaitForSingleObject(self.hWaitStop, win32event.INFINITE)

            event.clear()
            thread.join()

else:
    import logging
    log = logging.getLogger(__name__)
    def log_logger(*args):
        #log.error(_tb(args))
        log.error(*args)
    log_error = log_logger

def dummy_module(name, *objs):
    import imp
    mod = None
    parts = name.split('.')
    for i in range(len(parts)):
        part = '.'.join(parts[:i+1])
        sub = sys.modules[part] = imp.new_module(part)
        if mod is not None:
            setattr(mod, parts[i], sub)
        mod = sub
    for obj in objs:
        setattr(mod, obj.__name__, obj)
        obj.__module__ = name
    return mod


name = 'fenton.insecure'
if __name__ != name and name not in sys.modules:
    # make the error classes picklable under 'fenton.insecure'
    dummy_module(name, error, strerror, winerror)

if __name__ == "__main__":
    args = sys.argv[1:]
    if HAVE_WIN32:
        if not args or args[0] != 'test':
            win32serviceutil.HandleCommandLine(AppService)
            sys.exit()

        args.pop(0)

    if args:
        sec = args[0]
    else:
        sec = None
    test_server(sec)

